#include "detector/yolo11_onnx_detector.h"
#include <opencv2/imgproc.hpp>

YOLOv11Detector::YOLOv11Detector(const std::string& model_path)
    : OnnxDetector(model_path) {}

// YOLOv11输出假定为 [1, N, 85]，每行[x, y, w, h, obj_conf, cls0, cls1, ...]
#include <opencv2/dnn.hpp>
std::vector<DetectionResult> YOLOv11Detector::detect(const cv::Mat& image) {
    if (!initialized_) return {};
    // 1. 预处理与推理
    cv::Mat resized;
    cv::resize(image, resized, cv::Size(input_width_, input_height_));
    cv::Mat blob;
    resized.convertTo(blob, CV_32F, 1.0 / 255);
    std::vector<float> input_tensor_values(blob.total() * blob.channels());
    std::vector<cv::Mat> chw(blob.channels());
    for (int i = 0; i < blob.channels(); ++i)
        chw[i] = cv::Mat(blob.rows, blob.cols, CV_32F, input_tensor_values.data() + i * blob.rows * blob.cols);
    cv::split(blob, chw);
    std::array<int64_t, 4> input_shape = {1, blob.channels(), blob.rows, blob.cols};
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info, input_tensor_values.data(), input_tensor_values.size(), input_shape.data(), input_shape.size());
    const char* input_names_c[] = {input_name_.c_str()};
    const char* output_names_c[] = {output_name_.c_str()};
    auto output_tensors = session_->Run(Ort::RunOptions{nullptr}, 
        input_names_c, &input_tensor, 1, output_names_c, 1);
    // 2. 解析输出
    float conf_thresh = 0.4f, nms_thresh = 0.45f;
    std::vector<DetectionResult> results;
    auto& out_tensor = output_tensors[0];
    float* data = out_tensor.GetTensorMutableData<float>();
    auto shape = out_tensor.GetTensorTypeAndShapeInfo().GetShape();

    // reshape [1,84,8400] -> [1,8400,84] (转置后逐行访问)
    std::vector<float> transposed_data(shape[1] * shape[2]);
    for (int i = 0; i < shape[1]; ++i) {
        for (int j = 0; j < shape[2]; ++j) {
            transposed_data[j * shape[1] + i] = data[i * shape[2] + j];
        }
    }
    // 解析时以box为行（每行84维）

    /**
    std::cout << "[YOLOv11] Output shape: ";
    for (auto v : shape) std::cout << v << " ";
    std::cout << std::endl;
    if (shape.size() < 3) {
      	std::cerr << "[YOLOv11] Output shape error!" << std::endl;
        return {};
    }
    **/

    // output shape is [1,84,8400]
    int num_boxes = shape[2];
    int num_classes = shape[1] - 4;
    //std::cout << "[YOLOv11] num_boxes: " << num_boxes << ", num_classes: " << num_classes << std::endl;

    int stride = shape[1];
    int orig_w = image.cols, orig_h = image.rows;
    float scale_x = float(orig_w) / input_width_, scale_y = float(orig_h) / input_height_;
    std::vector<cv::Rect> boxes;
    std::vector<float> scores;
    std::vector<int> class_ids;
    for (int i = 0; i < num_boxes; ++i) {
        // [cx,cy,w,h, class0, class1, ..., classN]，此时每行84维
        float cx = transposed_data[i * stride + 0] * scale_x;
        float cy = transposed_data[i * stride + 1] * scale_y;
        float w  = transposed_data[i * stride + 2] * scale_x;
        float h  = transposed_data[i * stride + 3] * scale_y;
        // 取最大类别分数及其类别
        float max_score = -1;
        int max_cls = -1;
        for (int c = 0; c < num_classes; ++c) {
            float score = transposed_data[i * stride + (4 + c)];
            if (score > max_score) {
                max_score = score;
                max_cls = c;
            }
        }
        if (max_score < conf_thresh) continue;
        int left = int(cx - w / 2);
        int top  = int(cy - h / 2);
        boxes.emplace_back(left, top, int(w), int(h));
        scores.push_back(max_score);
        class_ids.push_back(max_cls);
    }
    std::vector<int> nms_indices;
    cv::dnn::NMSBoxes(boxes, scores, conf_thresh, nms_thresh, nms_indices);
    for (int idx : nms_indices) {
        DetectionResult dr;
        dr.bbox = boxes[idx] & cv::Rect(0, 0, orig_w, orig_h);
        dr.confidence = scores[idx];
        dr.class_id = class_ids[idx];
        results.push_back(dr);
    }
    return results;
}

BaseDetector* BaseDetector::CreateDetector(const std::string& model_path, int devId, int clusterId) {
    return new YOLOv11Detector(model_path, devId, clusterId);
}
